import os,re
import csv
from PIL import Image
import pytesseract

# 设置 Tesseract 可执行文件路径（Windows 用户需要设置）
# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

# 图片目录
image_dir = 'D:/Desk/0312-炳哥尺寸转换/尺寸标注/标注信息'

# 输出 CSV 文件路径
csv_output = 'D:/Desk/0312-炳哥尺寸转换/尺寸标注/results2.csv'

# 支持的图片格式
supported_formats = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')

# 获取所有图片路径
image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir)
               if f.lower().endswith(supported_formats)]

# 创建并写入 CSV 文件
with open(csv_output, 'w', newline='', encoding='utf-8-sig') as csvfile:
    # 定义 CSV 列名
    fieldnames = ['文件名', 'OCR结果']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    for idx, img_path in enumerate(image_paths, start=1):
        try:
            # 打开图片
            image = Image.open(img_path)
            
            # OCR 识别（中文）
            text = pytesseract.image_to_string(image, lang='chi_sim').strip()

            # 提取所有数字（包括整数和浮点数）
            numbers = re.findall(r'\d+\.?\d*', text)

            # 将字符串转换为浮点数或整数
            result = [float(num) if '.' in num else int(num) for num in numbers]
            
            # 写入 CSV 行
            writer.writerow({
                '文件名': os.path.basename(img_path),
                'OCR结果': result
            })

            # 打印进度（每处理100张提示一次）
            if idx % 10 == 0:
                print(f"已处理 {idx} 张图片")

        except Exception as e:
            # 记录错误信息到 CSV
            writer.writerow({
                '文件名': os.path.basename(img_path),
                'OCR结果': f'Error: {str(e)}'
            })

print(f"处理完成！结果已保存到 {csv_output}")